import os
import asyncio
import subprocess
import typer
import shutil
from metagpt.logs import logger
from metagpt.team import Team
from metagpt.schema import Message
from c2rust.roles import baseline_team, rag_team
from c2rust.roles.baseline import CompilerMessageAnalyser
from c2rust.actions import AnalyzeCode, AddMainFunc
import gradio as gr


def hire_team(use_rag=True, max_iters=5):
    def add_role(role_cls):
        if role_cls == CompilerMessageAnalyser:
            return CompilerMessageAnalyser(epoch_limit=max_iters)
        else:
            return role_cls()
    
    team_roles = []
    team_blueprint = None
    if use_rag:
        team_blueprint = rag_team
    else:
        team_blueprint = baseline_team
    
    for role_blueprint in team_blueprint:
        team_roles.append(add_role(role_blueprint))
    return team_roles
    

async def c2rust_run(
    c_code,
    investment = 3.0,
    max_round = 50
):
    # logger.info(c_code)
    investment = float(investment)
    n_round = int(max_round)
    
    # remove `workspace` dir, if exists
    if os.path.exists("workspace"):
        shutil.rmtree("workspace")

    team = Team()
    team.hire(hire_team(use_rag=True, max_iters=5))
    team.invest(investment=investment)
    
    # team.env.publish_message(
    #     Message(role="Human", content='''main.cpp;main.rs''', cause_by=AddMainFunc, send_to='<all>'),
    # )
    # await team.env.run(k=10)
    
    await team.run(n_round=n_round, idea=c_code)
    
    result = subprocess.run(["rustc", "main.rs"], capture_output=True, text=True)
    code_result = result.stderr
    # 返回 main.rs 中的内容
    with open('main.rs', 'r') as f:
        return (f.read(), code_result)
